# app/core/sql_generator.py
import re
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from typing import Dict, Any, Optional, Tuple
import uuid

class SQLGenerator:
    def __init__(self, user_id: str):
        """
        初始化SQL生成器
        :param user_id: 当前用户ID
        """
        self.user_id = user_id
        self.bill_table = "bill_info"
        self.type_table = "income_expenditure_type"
        self.user_table = "user_info"

        # 时间范围映射（与NLP模块一致）
        self.time_range_mapping = {
            "今天": ("today", 0),
            "昨天": ("yesterday", 1),
            "本周": ("this_week", 0),
            "上周": ("last_week", 1),
            "本月": ("this_month", 0),
            "上个月": ("last_month", 1),
            "今年": ("this_year", 0),
            "去年": ("last_year", 1)
        }

    def generate_sql(self, intent: str, entities: Dict[str, Any]) -> Optional[str]:
        """
        根据NLP解析结果生成SQL查询
        :param intent: NLP解析的意图
        :param entities: NLP提取的实体
        :return: 生成的SQL语句
        """
        # 映射NLP意图到SQL生成函数
        intent_mapping = {
            "add_expense": self._generate_insert_transaction,
            "query_expense": self._generate_query_transactions,
            "generate_report": self._generate_report_query,
            "financial_advice": self._generate_advice_query
        }

        if intent not in intent_mapping:
            return None

        return intent_mapping[intent](entities)

    def _generate_insert_transaction(self, entities: Dict[str, Any]) -> str:
        """
        生成添加交易记录的SQL
        """
        # 获取当前日期时间
        current_datetime = datetime.now()

        # 获取类别ID
        category_id = self._get_category_id(entities.get("category"))
        if not category_id:
            return ""

        # 获取收款用户ID
        get_user_id = self._get_user_id(entities.get("get_user_id"))

        # 构建插入语句
        columns = ["bill_id","time", "type_id", "money", "user_id", "note", "is_income"]

        values = [
            f"'{self.generate_bill_id()}'",
            f"'{self._format_datetime(current_datetime)}'",
            str(category_id),
            str(entities.get("money", 0)),
            f"'{str(self.user_id)}'",
            f"'{self.sanitize_input(entities.get('note', ''))}'",
            "1" if entities.get("is_income", False) else "0"
        ]

        # 添加收款用户ID
        if get_user_id:
            columns.append("get_user_id")
            values.append(str(get_user_id))

        sql = f"""
        INSERT INTO {self.bill_table} ({', '.join(columns)})
        VALUES ({', '.join(values)})
        """
        print("\n添加交易SQL:\n", sql)

        return sql

    def _generate_query_transactions(self, entities: Dict[str, Any]) -> str:
        """
        生成交易查询SQL
        """
        conditions = [f"b.user_id = '{self.user_id}'"]
        order = "ORDER BY b.time DESC"

        # 处理时间范围
        date_range = self._process_date_range(entities.get("date_range"))
        if date_range:
            conditions.append(f"b.time BETWEEN '{date_range[0]}' AND '{date_range[1]}'")

        # 金额过滤
        if "amount_range" in entities:
            min_amt, max_amt = entities["amount_range"]
            conditions.append(f"b.money BETWEEN {min_amt} AND {max_amt}")
        elif "amount" in entities:
            if entities["amount"] is not None:
                conditions.append(f"b.money = {entities['amount']}")

        # 分类过滤
        if "category" in entities:
            category_id = self._get_category_id(entities["category"])
            if category_id:
                conditions.append(f"b.type_id = {category_id}")

        # 收入/支出过滤
        # if "is_income" in entities:
        #     is_income = "1" if entities["is_income"] else "0"
        #     conditions.append(f"b.is_income = {is_income}")

        # 收款人过滤
        if "get_user" in entities:
            get_user_id = self._get_user_id(entities["get_user"])
            if get_user_id:
                conditions.append(f"b.get_user_id = '{get_user_id}'")
                print("\nconditions:\n", get_user_id)

        where_clause = " AND ".join(conditions)
        return f"""
        SELECT 
            b.bill_id, 
            b.time, 
            t.type_name AS category,
            b.money,
            u1.user_name AS payer,
            u2.user_name AS receiver,
            b.note,
            CASE b.is_income 
                WHEN 1 THEN '收入' 
                ELSE '支出' 
            END AS type
        FROM {self.bill_table} b
        JOIN {self.type_table} t ON b.type_id = t.type_id
        LEFT JOIN {self.user_table} u1 ON b.user_id = u1.user_id
        LEFT JOIN {self.user_table} u2 ON b.get_user_id = u2.user_id
        WHERE {where_clause}
        {order}
        """

    def _generate_report_query(self, entities: Dict[str, Any]) -> str:
        """
        生成财务报告SQL
        """
        # 确定报告类型
        report_type = entities.get("report_type", "summary")

        # 处理时间范围
        date_range = self._process_date_range(entities.get("date_range"))
        if not date_range:
            # 默认获取本月数据
            date_range = self._get_date_range("this_month")

        # 根据报告类型生成不同的SQL
        if report_type == "category_summary":
            return self._generate_category_summary(date_range)
        elif report_type == "trend":
            return self._generate_trend_query(date_range)
        else:
            return self._generate_summary_query(date_range)

    def _generate_summary_query(self, date_range: tuple) -> str:
        """
        生成财务摘要SQL
        """
        return f"""
        SELECT 
            SUM(CASE WHEN is_income = 1 THEN money ELSE 0 END) AS total_income,
            SUM(CASE WHEN is_income = 0 THEN money ELSE 0 END) AS total_expense,
            COUNT(*) AS transaction_count
        FROM {self.bill_table}
        WHERE user_id = {self.user_id}
          AND time BETWEEN '{date_range[0]}' AND '{date_range[1]}'
        """

    def _generate_category_summary(self, date_range: tuple) -> str:
        """
        生成分类统计SQL
        """
        return f"""
        SELECT 
            t.type_name AS category,
            SUM(b.money) AS total_amount,
            COUNT(*) AS transaction_count
        FROM {self.bill_table} b
        JOIN {self.type_table} t ON b.type_id = t.type_id
        WHERE b.user_id = {self.user_id}
          AND b.is_income = 0
          AND b.time BETWEEN '{date_range[0]}' AND '{date_range[1]}'
        GROUP BY t.type_name
        ORDER BY total_amount DESC
        """

    def _generate_trend_query(self, date_range: tuple) -> str:
        """
        生成消费趋势SQL（按月分组）
        """
        return f"""
        SELECT 
            DATE_FORMAT(b.time, '%Y-%m') AS month,
            SUM(b.money) AS total_expense
        FROM {self.bill_table} b
        WHERE b.user_id = {self.user_id}
          AND b.is_income = 0
          AND b.time BETWEEN '{date_range[0]}' AND '{date_range[1]}'
        GROUP BY month
        ORDER BY month
        """

    def _generate_advice_query(self, entities: Dict[str, Any]) -> str:
        """
        生成财务建议所需的SQL（获取历史数据）
        """
        # 默认获取过去6个月的数据
        end_date = datetime.now().strftime("%Y-%m-%d")
        start_date = (datetime.now() - relativedelta(months=6)).strftime("%Y-%m-%d")

        return f"""
        SELECT 
            t.type_name AS category, 
            b.money, 
            b.time
        FROM {self.bill_table} b
        JOIN {self.type_table} t ON b.type_id = t.type_id
        WHERE b.user_id = {self.user_id}
          AND b.is_income = 0
          AND b.time BETWEEN '{start_date}' AND '{end_date}'
        ORDER BY b.time DESC
        """

    def _get_category_id(self, category_name: str) -> Optional[int]:
        """
        获取类别ID
        """
        if not category_name:
            return None

        # 在实际应用中，这里应该查询数据库获取类别ID
        # 为简化演示，我们使用硬编码映射
        category_mapping = {
            "相互转账": 1,
            "早午晚餐": 2,
            "交通": 3,
            "零食水果": 4,
            "房租房贷": 5,
            "服饰鞋包": 6,
            "家居百货": 7,
            "礼金红包": 8,
            "医疗用品": 9,
            "护肤美容": 10,
            "话费网费": 11,
            "娱乐游玩": 12,
            "其他": 13,
            "零碎未记": 14,
            "工资收入": 15,
            "大额支出": 16,
            "彤宝用品": 17,
            "人情往来": 18,
            "水电燃气": 19
        }

        return category_mapping.get(category_name)

    def _get_user_id(self, user_name: str) -> Optional[int]:
        """
        获取用户ID
        """
        if not user_name:
            return None

        # 在实际应用中，这里应该查询数据库获取用户ID
        # 为简化演示，我们使用硬编码映射
        user_mapping = {
            "张三": 101,
            "李四": 102,
            "王五": 103
        }

        return user_mapping.get(user_name)

    def _process_date_range(self, date_range: Any) -> Optional[tuple]:
        """
        处理时间范围实体，转换为实际的日期范围
        """
        if not date_range:
            return None

        # 如果已经是元组格式（YYYY-MM-DD）
        if isinstance(date_range, tuple) and len(date_range) == 2:
            return date_range

        # 如果是预定义的时间范围
        if date_range in self.time_range_mapping:
            return self._get_date_range(date_range)

        return None

    def _get_date_range(self, time_range: str) -> tuple:
        """
        根据预定义的时间范围获取实际的日期范围
        """
        now = datetime.now()
        start = ""
        end = self._format_datetime(now)

        if time_range == "today":
            start = self._format_datetime(now.replace(hour=0, minute=0, second=0, microsecond=0))
        elif time_range == "yesterday":
            yesterday = now - timedelta(days=1)
            start = self._format_datetime(yesterday.replace(hour=0, minute=0, second=0, microsecond=0))
            end = self._format_datetime(yesterday.replace(hour=23, minute=59, second=59, microsecond=999999))
        elif time_range == "this_week":
            # 获取本周一的日期
            monday = now - timedelta(days=now.weekday())
            start = self._format_datetime(monday.replace(hour=0, minute=0, second=0, microsecond=0))
        elif time_range == "last_week":
            # 获取上周日的日期
            last_sunday = now - timedelta(days=now.weekday() + 1)
            # 获取上周一的日期
            last_monday = last_sunday - timedelta(days=6)

            start = self._format_datetime(last_monday.replace(hour=0, minute=0, second=0, microsecond=0))
            end = self._format_datetime(last_sunday.replace(hour=23, minute=59, second=59, microsecond=999999))
        elif time_range == "this_month":
            # 获取本月第一天
            first_day = now.replace(day=1)
            start = self._format_datetime(first_day.replace(hour=0, minute=0, second=0, microsecond=0))
        elif time_range == "last_month":
            # 获取上个月最后一天
            first_day = now.replace(day=1)
            last_month_last_day = first_day - timedelta(days=1)

            # 获取上个月第一天
            last_month_first_day = last_month_last_day.replace(day=1)

            start = self._format_datetime(last_month_first_day.replace(hour=0, minute=0, second=0, microsecond=0))
            end = self._format_datetime(last_month_last_day.replace(hour=23, minute=59, second=59, microsecond=999999))
        elif time_range == "this_year":
            # 获取今年第一天
            first_day = now.replace(month=1, day=1)
            start = self._format_datetime(first_day.replace(hour=0, minute=0, second=0, microsecond=0))
        elif time_range == "last_year":
            # 获取去年最后一天
            last_year_last_day = now.replace(year=now.year - 1, month=12, day=31)

            # 获取去年第一天
            last_year_first_day = now.replace(year=now.year - 1, month=1, day=1)

            start = self._format_datetime(last_year_first_day.replace(hour=0, minute=0, second=0, microsecond=0))
            end = self._format_datetime(last_year_last_day.replace(hour=23, minute=59, second=59, microsecond=999999))
        else:
            # 默认返回本月
            first_day = now.replace(day=1)
            start = self._format_datetime(first_day.replace(hour=0, minute=0, second=0, microsecond=0))

        return (start, end)

    def _format_datetime(self, dt: datetime) -> str:
        """
        将datetime对象格式化为2025072001315格式
        """
        date_str = dt.strftime("%Y%m%d")
        weekday = dt.weekday()  # 0-6，周一到周日
        time_str = dt.strftime("%H%M")
        return f"{date_str}{weekday}{time_str}"

    @staticmethod
    def sanitize_input(input_str: str) -> str:
        """
        防止SQL注入的安全处理
        """
        return input_str if input_str is None else re.sub(r"[;'\"]", "", input_str)

    def generate_bill_id(self) -> str:
        # 生成UUID4（随机生成）并转换为十六进制字符串（不含连字符）
        return uuid.uuid4().hex
# 使用示例
if __name__ == "__main__":
    generator = SQLGenerator(user_id='cai')

    # 示例1：添加支出
    sql = generator.generate_sql(
        intent="add_expense",
        # entities={
        #     "amount": 150.0,
        #     "category": "早午晚餐",
        #     "is_income": False,
        #     "note": "午餐外卖"
        # }
        entities={'money': 35, 'category': '早午晚餐', 'date_range': '昨天', 'merchant': '星巴克', 'is_income': False,  'report_type': None,  "note": None}
    )
    print("添加支出SQL:\n", sql)

    # 示例2：添加收入
    # sql = generator.generate_sql(
    #     intent="add_expense",
    #     entities={
    #         "amount": 8000.0,
    #         "category": "工资收入",
    #         "is_income": True,
    #         "note": "8月工资"
    #     }
    # )
    # print("\n添加收入SQL:\n", sql)

    # # 示例3：添加转账
    # sql = generator.generate_sql(
    #     intent="add_expense",
    #     entities={
    #         "amount": 500.0,
    #         "category": "相互转账",
    #         "is_income": False,
    #         "get_user_id": "李四",
    #         "note": "借款给李四"
    #     }
    # )
    # print("\n添加转账SQL:\n", sql)
    #
    # # 示例4：查询本月餐饮支出
    # sql = generator.generate_sql(
    #     intent="query_expense",
    #     entities={
    #         "category": "早午晚餐",
    #         "date_range": "本月",
    #         "is_income": False
    #     }
    # )
    # print("\n查询餐饮支出SQL:\n", sql)
    #
    # # 示例5：生成分类报告
    # sql = generator.generate_sql(
    #     intent="generate_report",
    #     entities={
    #         "report_type": "category_summary",
    #         "date_range": "本月"
    #     }
    # )
    # print("\n分类报告SQL:\n", sql)